from solvers.solverbase import SolverBase
from solvers.codina.codinabase import CodinaBase
from solvers.stabilizations import supg
from dolfin import *
import numpy

class GaussLegendre4(SolverBase,CodinaBase):
    '2-symplectic method rewritten without increments. Nonlinear-crosswind stabilized.'
    def __init__(self,p=1):
        SolverBase.__init__(self,p)
        CodinaBase.__init__(self)

    def solve(self,problem):
        mesh = problem.domain

        V = FunctionSpace(mesh,'CG',1)
        VV = VectorFunctionSpace(mesh,'CG',1)
        DG = FunctionSpace(mesh,'DG',0)

        def get_sigma(v,u_,u__,eps,dt,sigma,cache): # u_ is the solution @ current iteration, u__ is a solution of previous time level
            phiS = TestFunction(DG)  # universal test function   
            trial = TrialFunction(DG) # trial function for seminorm and norm of solution and residuum 
        
            solSemF = Function(DG) #this is where the values will go
            resF = Function(DG)
            solNormF = Function(DG)            
      
            if 'a' not in cache.keys():
                a = trial*phiS*dx                                             # a
                cache['a'] = assemble(a)
        #----
            solSeml = sqrt(inner(nabla_grad(u_),nabla_grad(u_)))*phiS*dx   # l1 
            b1 = assemble(solSeml)
            solve(cache['a'],solSemF.vector(),b1)
        #-----     
            resl = (u_-u__+dt*inner(v,nabla_grad(u_)))*phiS*dx                 # l2               #TODO consider adding time derivate
            b2 = assemble(resl)
            solve(cache['a'],resF.vector(),b2)
        #-----
            solNorml = sqrt(resF*resF)*phiS*dx                                   # l3
            b3 = assemble(solNorml)
            solve(cache['a'],solNormF.vector(),b3)
        #-----
            
            nc = mesh.num_cells()
            sigmaValues = numpy.zeros(nc)
            beta = 0.7
            for i in xrange(nc): # build the values
                c = Cell(mesh,i)
               
                diamK = c.diameter()  #elements diameter
                x = c.midpoint()
                
                norm = solNormF(x.x(),x.y())
                seminorm = solSemF(x.x(),x.y())
                
                if norm < 1E-16 or seminorm < 1E-16:
                    sigmaValues[i] = 0.
                else:
                    sigmaValues[i] = 0.5*max(0,beta - 2*eps*seminorm/norm/diamK)*diamK*norm/seminorm
                
            sigma.vector()[:] = sigmaValues
        #-----------------------------------------------------------------

        phi_ = interpolate(problem.phi_,V)
        phiP = Function(V)                  # solution to the auxiliary equation
        phiC = Function(V)                  # solution to the true equation
        phi = TrialFunction(V)          
        varphi = TestFunction(V)

        t = problem.t          
        dt = Constant(SolverBase._get_time_step(self,problem,VV,t))
        T = problem.T

        v = SolverBase._get_velocity(self,problem,VV,t)
        
        tau = Function(DG)
        tau.vector()[:] = supg.get_tau(v,supg.diffusivity)  # stab. parameters

        sigma_p = Function(DG) 
        sigma_m = Function(DG)
        pi_p = Function(DG)
        pi_m = Function(DG) 
        
        pert = conditional(
                           gt(sqrt(inner(v,v)),1E-16),
                           Constant(1.0),
                           Constant(0.) 
                           ) # a way to realize swith on |v| == 0.

        pMatrix = conditional(
                             gt(sqrt(inner(v,v)),1E-16),
                             Identity(2)-outer(v,v)/inner(v,v),
                             Constant(0.)*Identity(2),
                             ) # I - v X v
 
        aP = phi*varphi*dx + tau*pert*phi*inner(v,nabla_grad(varphi))*dx\
             +0.5*dt*sqrt(3.)/6.*sigma_m*inner(dot(nabla_grad(phi),pMatrix),nabla_grad(varphi))*dx\
             -0.5*dt*sqrt(3.)/6.*sigma_p*inner(dot(nabla_grad(phi),pMatrix),nabla_grad(varphi))*dx

        lP = phi_*varphi*dx\
             + dt*sqrt(3.)/6.*inner(v,nabla_grad(phiC))*varphi*dx\
             - dt*sqrt(3.)/6.*inner(v,nabla_grad(phi_))*varphi*dx\
             + tau*pert*phi_*inner(v,nabla_grad(varphi))*dx \
             + dt*sqrt(3.)/6.*tau*pert*inner(dot(nabla_grad(phiC),outer(v,v)),nabla_grad(varphi))*dx \
             - dt*sqrt(3.)/6.*tau*pert*inner(dot(nabla_grad(phi_),outer(v,v)),nabla_grad(varphi))*dx\
             -0.5*dt*1./2.*sigma_m*inner(dot(nabla_grad(phi_),pMatrix),nabla_grad(varphi))*dx\
             -0.5*dt*(1./2.-sqrt(3.)/6.)*sigma_m*inner(dot(nabla_grad(phiC),pMatrix),nabla_grad(varphi))*dx\
             +0.5*dt*(1./2.-2.*sqrt(3.)/6.)*sigma_p*inner(dot(nabla_grad(phi_),pMatrix),nabla_grad(varphi))*dx\
             +0.5*dt*(1./2.+sqrt(3.)/6.)*sigma_p*inner(dot(nabla_grad(phiC),pMatrix),nabla_grad(varphi))*dx
#-----
        aC = phi*varphi*dx + 0.5*dt*inner(v,nabla_grad(phi))*varphi*dx\
             +tau*pert*phi*inner(v,nabla_grad(varphi))*dx\
             +0.5*dt*tau*pert*inner(dot(nabla_grad(phi),outer(v,v)),nabla_grad(varphi))*dx\
              +0.5*dt*(1./2.-sqrt(3.)/6.)*pi_m*inner(dot(nabla_grad(phi),pMatrix),nabla_grad(varphi))*dx\
              +0.5*dt*(1./2.+sqrt(3.)/6.)*pi_p*inner(dot(nabla_grad(phi),pMatrix),nabla_grad(varphi))*dx
    
        lC = phi_*varphi*dx + tau*pert*phi_*inner(v,nabla_grad(varphi))*dx\
             -dt*(0.5-sqrt(3.)/6.)*inner(v,nabla_grad(phi_))*varphi*dx\
             -dt*(0.5-sqrt(3.)/6.)*tau*pert*inner(dot(nabla_grad(phi_),outer(v,v)),nabla_grad(varphi))*dx\
             -dt*sqrt(3.)/6.*inner(v,nabla_grad(phiP))*varphi*dx\
             -dt*sqrt(3.)/6.*tau*pert*inner(dot(nabla_grad(phiP),outer(v,v)),nabla_grad(varphi))*dx\
             -0.5*dt*1./2.*pi_m*inner(dot(nabla_grad(phi_),pMatrix),nabla_grad(varphi))*dx\
             -0.5*dt*sqrt(3.)/6.*pi_m*inner(dot(nabla_grad(phiP),pMatrix),nabla_grad(varphi))*dx\
             -0.5*dt*(1./2.-2.*sqrt(3.)/6.)*pi_p*inner(dot(nabla_grad(phi_),pMatrix),nabla_grad(varphi))*dx\
             -0.5*dt*sqrt(3.)/6.*pi_p*inner(dot(nabla_grad(phiP),pMatrix),nabla_grad(varphi))*dx

        self.update(problem,phi_,t,float(dt),0)  # store the solution at time 0

        phiC.assign(phi_)   # initial guess
        
        phiC_ = Function(V) # so that error could be calculated
        phiC_.assign(phi_)

        phi = Function(V)
        phi_k = Function(V) # iterations for 1st eq
        phi_m = Function(V) # iterations for 2nd eq
   
        t = 0

        count = 0
        cache = {}
        maxIter = 10
        tolG = 1E-8            # tolerance for gauss loop
        tolP = 1E-3            # tolerance for picard loop
        while t <= T:  # time loop
            t += float(dt)
            count += 1
            
            errorG = 1.
            iterationG = 0
           
            while errorG > tolG and iterationG < maxIter: # gauss4 loop
                iterationG += 1

                errorP1 = 1.
                iterationP1 = 0
                # guess for sigma_m 
                get_sigma(v,(0.5*phi_+(0.5-sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phi_),phi_,supg.diffusivity,float(dt),sigma_m,cache) 
                # guess for sigma_p
                get_sigma(v,((0.5-2*sqrt(3.)/6.)*phi_+(0.5+sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phi_),phi_,supg.diffusivity,float(dt),sigma_p,cache)  
                while errorP1 > tolP and iterationP1 < maxIter: # picard loop for second eq.
                    iterationP1 += 1
                    AP = assemble(aP)
                    LP = assemble(lP)
                    solve(AP,phi.vector(),LP)  # solve true system
                    diffP1 = phi.vector().array() - phi_k.vector().array()
                    errorP1 = numpy.linalg.norm(diffP1,ord=numpy.Inf)
                    phi_k.assign(phi)
                    # guess for sigma_m 
                    get_sigma(v,(0.5*phi_+(0.5-sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phi_k),phi_,supg.diffusivity,float(dt),sigma_m,cache) 
                    # guess for sigma_p
                    get_sigma(v,((0.5-2*sqrt(3.)/6.)*phi_+(0.5+sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phi_k),phi_,supg.diffusivity,float(dt),sigma_p,cache)  
                    #print '\tpicard1: ',iterationP1,errorP1
                phiP.assign(phi)

                errorP2 = 1.
                iterationP2 = 0
                # guess for pi_m 
                get_sigma(v,(0.5*phi_+(0.5-sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phiP),phi_,supg.diffusivity,float(dt),pi_m,cache) 
                # guess for pi_p
                get_sigma(v,((0.5-2*sqrt(3.)/6.)*phi_+(0.5+sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phiP),phi_,supg.diffusivity,float(dt),pi_p,cache)  
                while errorP2 > tolP and iterationP2 < maxIter: # picard loop for second eq.
                    iterationP2 += 1
                    AC = assemble(aC)
                    LC = assemble(lC)
                    solve(AC,phi.vector(),LC)  # solve true system
                    diffP2 = phi.vector().array() - phi_m.vector().array()
                    errorP2 = numpy.linalg.norm(diffP2,ord=numpy.Inf)
                    phi_m.assign(phi)
                    # guess for pi_m 
                    get_sigma(v,(0.5*phi_+(0.5-sqrt(3.)/6.)*phi_m+sqrt(3.)/6.*phiP),phi_,supg.diffusivity,float(dt),pi_m,cache) 
                    # guess for pi_p
                    get_sigma(v,((0.5-2*sqrt(3.)/6.)*phi_+(0.5+sqrt(3.)/6.)*phi_m+sqrt(3.)/6.*phiP),phi_,supg.diffusivity,float(dt),pi_p,cache)  
                    #print '\tpicard2: ',iterationP2,errorP2
                phiC.assign(phi)

                diffG = phiC.vector().array() - phiC_.vector().array()
                errorG = numpy.linalg.norm(diffG,ord=2)        # compute the error
                phiC_.assign(phiC)
                #print 'gauss: ',iterationG,errorG

            phi_.assign(phiC)
            self.update(problem,phi_,t,float(dt),count)

        problem.save_data(self.saveDir,str(self.CFL))
